import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

import torch,os,pdb
from diffusers import FluxPriorReduxPipeline, FluxPipeline,FluxFillPipeline
# from diffusers.utils import load_image
from util_flux import horizontal_concat_images
from util_flux import process_img_1024,concat_half
from image_gen_aux import DepthPreprocessor
from MODEL_CKP import DEPTH_PREDCITION

examples_dir = '/mnt/nas/shengjie/datasets/cloth_collar_localimg'
save_dir = '/data/shengjie/synthesis1/'

# imagefiles = os.listdir(examples_dir)



FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_ADAPTER = '/home/shengjie/ckp/flux-ip-adapter-v2'
FLUX_ADAPTER_ENCODER = '/home/shengjie/ckp/clip-vit-large-patch14'
FLUX='/data/models/FLUX___1-dev'

processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)


pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                    FLUX_REDUX, 
                                    torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxPipeline.from_pretrained(
    FLUX , 
    text_encoder=None,
    text_encoder_2=None,
    torch_dtype=torch.bfloat16
).to("cuda")
# pipe.load_ip_adapter(FLUX_ADAPTER,weight_name='ip_adapter.safetensors',
#                      image_encoder_pretrained_model_name_or_path=FLUX_ADAPTER_ENCODER)

# pdb.set_trace()

clo_list = []
for entry in os.scandir(examples_dir):
    filename = entry.name
    if not filename.endswith('.jpg'):continue
    # break
    clo_list.append(filename)
    if len(clo_list)==2:
        filename,filename2 = clo_list
        clo_list.clear()
        test_img = os.path.join(examples_dir,filename)
        test_img2 = os.path.join(examples_dir,filename2)
    else:
        continue
    # ip_img = os.path.join(examples_dir,imagefiles[1])



    test_steps = 8
    # test_steps = 10
    # save_test_img = os.path.join(save_dir,
    #                             os.path.splitext(imagefiles[0])[0]+f'_step={test_steps}'+\
    #                                 os.path.splitext(imagefiles[0])[1]
    #                             )
    # pdb.set_trace()

    # image = load_image( test_img )
    # ip_img = load_image( ip_img )

    target_shape = (1024,1024)
    image = process_img_1024(test_img)
    image2 = process_img_1024(test_img2)


    # dpth1 = processor( image )[0].convert('RGB')
    # dpth2 = processor( image2 )[0].convert('RGB')

    # img = image + image2
    # 取出 image 的上一半 + image2 的下一半
  
    # concat_type = input('input concat type [h,v,c] : ')
    # image = concat_half( test_img, test_img2,concat_type=concat_type )

    # pdb.set_trace()

    with torch.no_grad():
        prompt_embeds,pooled_prompt_embeds = \
            pipe_prior_redux(image,return_dict=False)
        prompt_embeds2,pooled_prompt_embeds2 = \
            pipe_prior_redux(image2,return_dict=False)

        alpha = 0.4
        prompt_emb = (1-alpha) * prompt_embeds + alpha * prompt_embeds2
        pooled_prompt_emb = (1-alpha) * pooled_prompt_embeds + alpha * pooled_prompt_embeds2


        images = pipe(
            guidance_scale=4.5,
            num_inference_steps=20,
            # generator=torch.Generator("cpu").manual_seed(0),
            # ip_adapter_image=ip_img,
            # **pipe_prior_output,
            prompt_embeds=prompt_emb,
            pooled_prompt_embeds=pooled_prompt_emb,
        ).images

    # print(len(images))
    concat_img = horizontal_concat_images([image,image2,images[0]])
    concat_img.save('./tmp.jpg')

    pdb.set_trace()

    # images[0].save('tmp_redux.jpg')